from enum import Enum
from ipaddress import ip_network, ip_address

from scapy.layers.inet import TCP, UDP, IP


class Action(Enum):
    ALERT = 1


def action(istr: str):  # alert
    if istr.lower().strip() == 'action':
        return Action.ALERT


class IPNetwork(object):  # ip地址
    def __init__(self, text):
        if text.strip() == 'any':
            self.ipn = ip_network(u'0.0.0.0/0')
        else:
            strs = text.strip().split('/')
            if len(strs) >= 2:
                self.ipn = ip_network(text.strip())
            else:
                self.ipn = ip_network(strs[0] + '/32')

    def contains(self, ip):
        return ip in self.ipn


class Ports:  # 端口
    def __init__(self, text):
        if text.strip() == 'any':
            self.type = 'any'
        elif ':' in text:
            self.type = 'range'
            strs = text.strip().split(':')
            self.low_port = int(strs[0])
            self.high_port = int(strs[1])
        elif ',' in text:
            self.type = 'list'
            self.list_ports = list()
            for s in text.strip().split(','):
                self.list_ports.append(int(s))
        else:
            self.type = 'list'
            self.list_ports = [int(text.strip()), ]

    def contains(self, port):
        if self.type == 'any':
            return True
        elif self.type == 'range':
            return self.low_port <= port <= self.high_port
        else:
            return port in self.list_ports


class Protocol(Enum):  # 协议
    TCP = 6
    UDP = 17


def protocol(istr: str):
    if istr.lower().strip() == 'tcp':
        return Protocol.TCP
    else:
        return Protocol.UDP


class IDSRule(object):
    def __init__(self, text):
        # 构建一条规则
        self.text = text
        strs = self.text.strip().split('\t')
        if len(strs) >= 7:
            self.action = action(strs[0])
            self.protocol = protocol(strs[1])
            self.src = IPNetwork(strs[2])
            self.sport = Ports(strs[3])
            self.dst = IPNetwork(strs[5])
            self.dport = Ports(strs[6])
            opts = strs[7].strip()[1:-1].strip().split(';')
            for opt in opts:
                if opt == '' or opt == ')' or len(opt.split(':')) != 2:
                    continue
                option, value = opt.split(':')
                if option == 'msg':
                    self.msg = value
                elif opts == 'tos':
                    self.tos = int(value)
                elif option == 'len':
                    self.len = int(value)
                elif option == 'offset':
                    self.offset = int(value)
                elif option == 'seq':
                    self.seq = int(value)
                elif option == 'ack':
                    self.ack = int(value)
                elif option == 'flags':
                    self.flags = value
                else:
                    self.default = value

    def check_protocol(self, proto: int):
        if self.protocol.value == proto:
            return True
        return False

    def check_ips(self, pkt: IP):
        src = pkt[IP].src
        dst = pkt[IP].dst
        ip_src = ip_address(src)
        ip_dst = ip_address(dst)
        return self.src.contains(ip_src) and self.dst.contains(ip_dst)

    def check_ports(self, pkt: IP, proto: int):
        if proto == 17:  # UDP
            sport = pkt[UDP].sport
            dport = pkt[UDP].dport
        else:
            sport = pkt[TCP].sport
            dport = pkt[TCP].dport
        return self.sport.contains(sport) and self.dport.contains(dport)

    def check_options(self, pkt: IP, proto):
        if hasattr(self, 'tos') and self.tos != int(pkt[IP].tos):
            return False

        if hasattr(self, 'len') and self.len != int(pkt[IP].len):
            return False

        if hasattr(self, 'offset') and self.offset != int(pkt[IP].frag):
            return False

        if proto == 6:  # TCP
            if hasattr(self, 'seq') and self.seq != int(pkt[TCP].seq):
                return False
            if hasattr(self, 'ack') and self.seq != int(pkt[TCP].seq):
                return False
        return True

    def match(self, pkt):  # True is not matched
        proto = pkt[IP].proto
        return self.check_protocol(proto) and self.check_ports(pkt, proto) and self.check_ips(
            pkt) and self.check_options(pkt, proto)
